#!/usr/bin/env python3
"""
安全帽检测推理脚本
使用训练好的YOLOv8模型进行安全帽检测
"""

import os
import cv2
import numpy as np
from ultralytics import YOLO
import matplotlib.pyplot as plt
import argparse
from pathlib import Path

class HelmetDetector:
    def __init__(self, model_path="runs/detect/helmet_detection/weights/best.pt"):
        """
        初始化安全帽检测器
        
        Args:
            model_path: 训练好的模型路径
        """
        self.model_path = model_path
        self.model = None
        self.class_names = ['Helmet', 'No-Helmet', 'Person']
        self.colors = {
            'Helmet': (0, 255, 0),      # 绿色 - 佩戴安全帽
            'No-Helmet': (0, 0, 255),   # 红色 - 未佩戴安全帽
            'Person': (255, 0, 0)       # 蓝色 - 人员
        }
        
        self.load_model()
    
    def load_model(self):
        """加载训练好的模型"""
        if not os.path.exists(self.model_path):
            print(f"错误: 找不到模型文件 {self.model_path}")
            print("请先运行 train_helmet_detection.py 训练模型")
            return False
        
        try:
            self.model = YOLO(self.model_path)
            print(f"成功加载模型: {self.model_path}")
            return True
        except Exception as e:
            print(f"加载模型失败: {e}")
            return False
    
    def detect_image(self, image_path, conf_threshold=0.5, save_result=True):
        """
        对单张图片进行检测
        
        Args:
            image_path: 图片路径
            conf_threshold: 置信度阈值
            save_result: 是否保存检测结果
        
        Returns:
            检测结果
        """
        if self.model is None:
            print("模型未加载")
            return None
        
        # 读取图片
        image = cv2.imread(image_path)
        if image is None:
            print(f"无法读取图片: {image_path}")
            return None
        
        # 进行检测
        results = self.model(image, conf=conf_threshold)
        
        # 绘制检测结果
        annotated_image = self.draw_results(image, results[0])
        
        if save_result:
            # 保存结果
            output_dir = "detection_results"
            os.makedirs(output_dir, exist_ok=True)
            
            filename = Path(image_path).stem
            output_path = os.path.join(output_dir, f"{filename}_detected.jpg")
            cv2.imwrite(output_path, annotated_image)
            print(f"检测结果已保存到: {output_path}")
        
        return results[0], annotated_image
    
    def detect_video(self, video_path, conf_threshold=0.5, save_result=True):
        """
        对视频进行检测
        
        Args:
            video_path: 视频路径
            conf_threshold: 置信度阈值
            save_result: 是否保存检测结果
        """
        if self.model is None:
            print("模型未加载")
            return
        
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            print(f"无法打开视频: {video_path}")
            return
        
        # 获取视频信息
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        
        # 设置输出视频
        if save_result:
            output_dir = "detection_results"
            os.makedirs(output_dir, exist_ok=True)
            
            filename = Path(video_path).stem
            output_path = os.path.join(output_dir, f"{filename}_detected.mp4")
            
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        
        frame_count = 0
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        print(f"开始处理视频: {video_path}")
        print(f"总帧数: {total_frames}, FPS: {fps}")
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            frame_count += 1
            
            # 进行检测
            results = self.model(frame, conf=conf_threshold)
            
            # 绘制检测结果
            annotated_frame = self.draw_results(frame, results[0])
            
            # 显示进度
            if frame_count % 30 == 0:
                progress = (frame_count / total_frames) * 100
                print(f"处理进度: {progress:.1f}% ({frame_count}/{total_frames})")
            
            if save_result:
                out.write(annotated_frame)
            
            # 显示实时结果(可选)
            cv2.imshow('Helmet Detection', annotated_frame)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        
        cap.release()
        if save_result:
            out.release()
            print(f"检测结果已保存到: {output_path}")
        
        cv2.destroyAllWindows()
    
    def draw_results(self, image, results):
        """
        在图片上绘制检测结果
        
        Args:
            image: 原始图片
            results: 检测结果
        
        Returns:
            标注后的图片
        """
        annotated_image = image.copy()
        
        if results.boxes is not None:
            for box in results.boxes:
                # 获取边界框坐标
                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
                
                # 获取类别和置信度
                conf = box.conf[0].cpu().numpy()
                cls = int(box.cls[0].cpu().numpy())
                class_name = self.class_names[cls]
                
                # 选择颜色
                color = self.colors.get(class_name, (128, 128, 128))
                
                # 绘制边界框
                cv2.rectangle(annotated_image, (x1, y1), (x2, y2), color, 2)
                
                # 绘制标签
                label = f"{class_name}: {conf:.2f}"
                label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
                
                # 标签背景
                cv2.rectangle(annotated_image, 
                            (x1, y1 - label_size[1] - 10), 
                            (x1 + label_size[0], y1), 
                            color, -1)
                
                # 标签文字
                cv2.putText(annotated_image, label, 
                          (x1, y1 - 5), 
                          cv2.FONT_HERSHEY_SIMPLEX, 0.6, 
                          (255, 255, 255), 2)
        
        return annotated_image
    
    def analyze_safety_compliance(self, results):
        """
        分析安全合规性
        
        Args:
            results: 检测结果
        
        Returns:
            安全分析报告
        """
        if results.boxes is None:
            return {
                'total_persons': 0,
                'helmet_wearing': 0,
                'no_helmet': 0,
                'compliance_rate': 0.0,
                'safety_status': 'No persons detected'
            }
        
        person_count = 0
        helmet_count = 0
        no_helmet_count = 0
        
        for box in results.boxes:
            cls = int(box.cls[0].cpu().numpy())
            class_name = self.class_names[cls]
            
            if class_name == 'Person':
                person_count += 1
            elif class_name == 'Helmet':
                helmet_count += 1
            elif class_name == 'No-Helmet':
                no_helmet_count += 1
        
        # 计算合规率
        total_head_detection = helmet_count + no_helmet_count
        compliance_rate = (helmet_count / total_head_detection * 100) if total_head_detection > 0 else 0
        
        # 安全状态评估
        if compliance_rate >= 90:
            safety_status = "优秀"
        elif compliance_rate >= 70:
            safety_status = "良好"
        elif compliance_rate >= 50:
            safety_status = "一般"
        else:
            safety_status = "需要改进"
        
        return {
            'total_persons': person_count,
            'helmet_wearing': helmet_count,
            'no_helmet': no_helmet_count,
            'compliance_rate': compliance_rate,
            'safety_status': safety_status
        }

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='安全帽检测')
    parser.add_argument('--input', '-i', required=True, help='输入图片或视频路径')
    parser.add_argument('--model', '-m', default='runs/detect/helmet_detection/weights/best.pt', 
                       help='模型路径')
    parser.add_argument('--conf', '-c', type=float, default=0.5, help='置信度阈值')
    parser.add_argument('--save', '-s', action='store_true', help='保存检测结果')
    
    args = parser.parse_args()
    
    # 创建检测器
    detector = HelmetDetector(args.model)
    
    input_path = args.input
    
    # 检查输入文件是否存在
    if not os.path.exists(input_path):
        print(f"错误: 找不到输入文件 {input_path}")
        return
    
    # 判断输入类型
    file_ext = Path(input_path).suffix.lower()
    image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']
    video_extensions = ['.mp4', '.avi', '.mov', '.mkv', '.flv']
    
    if file_ext in image_extensions:
        print(f"检测图片: {input_path}")
        results, annotated_image = detector.detect_image(input_path, args.conf, args.save)
        
        if results is not None:
            # 分析安全合规性
            safety_report = detector.analyze_safety_compliance(results)
            
            print("\n安全分析报告:")
            print(f"检测到人员: {safety_report['total_persons']}")
            print(f"佩戴安全帽: {safety_report['helmet_wearing']}")
            print(f"未佩戴安全帽: {safety_report['no_helmet']}")
            print(f"合规率: {safety_report['compliance_rate']:.1f}%")
            print(f"安全状态: {safety_report['safety_status']}")
            
            # 显示结果
            plt.figure(figsize=(12, 8))
            plt.imshow(cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB))
            plt.title('安全帽检测结果')
            plt.axis('off')
            plt.show()
    
    elif file_ext in video_extensions:
        print(f"检测视频: {input_path}")
        detector.detect_video(input_path, args.conf, args.save)
    
    else:
        print(f"不支持的文件格式: {file_ext}")
        print(f"支持的图片格式: {image_extensions}")
        print(f"支持的视频格式: {video_extensions}")

if __name__ == "__main__":
    main() 